"""
@Desc: Predicting Stock Prices with Linear Regression
@Auth: meihongliang-m2
@Date: 2025/3/18-16:15
"""

import yfinance as yf
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt

# Step 1: Fetch historical stock price data #1.获取历史股票价格数据
ticker = "AAPL"  # Example: Apple Inc.
data = yf.download(ticker, start="2015-01-01", end="2023-01-01")

# Step 2: Prepare the data for the model #2. 为模型准备数据
data['Return'] = data['Adj Close'].pct_change()  # Calculate daily returns
data['Target'] = data['Return'].shift(-1)  # Target is the next day's return
data.dropna(inplace=True)  # Remove rows with NaN values

# Step 3: Define features (X) and target (y)  #3.定义特征X和目标Y
X = data[['Return']]
y = data['Target']

# Step 4: Split the data into training and testing sets #4.将数据拆分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Step 5: Create and train the model #5.创建和训练模型
model = LinearRegression()
model.fit(X_train, y_train)

# Step 6: Make predictions #6.做出预测
predictions = model.predict(X_test)

# Step 7: Visualize the results #7.可视化结果
plt.figure(figsize=(10, 6))
plt.scatter(y_test, predictions)
plt.xlabel("Actual Returns")
plt.ylabel("Predicted Returns")
plt.title("Actual vs Predicted Returns")
plt.plot([-0.1, 0.1], [-0.1, 0.1], color='red')  # Perfect prediction line
plt.show()

# Step 8: Evaluate the model #8.评估模型
accuracy = model.score(X_test, y_test)
print(f"Model accuracy: {accuracy:.2f}")
